import torch
import torch.nn as nn
import torch.nn.functional as F

from .preprocess_utils import *


class EpipolarLoss_full(nn.Module):
    def __init__(self, configs, device=None):
        super(EpipolarLoss_full, self).__init__()
        self.__lossname__ = 'EpipolarLoss_fullinfo'
        self.config = configs
        self.w_g = self.config['weight_grid']
        self.w_w = self.config['weight_window']

    def epipolar_cost(self, coord1, coord2, fmatrix, im_size):
        coord1_h = homogenize(coord1).transpose(1, 2)
        coord2_h = homogenize(coord2).transpose(1, 2)
        epipolar_line = fmatrix.bmm(coord1_h)  # Bx3xn
        epipolar_line_ = epipolar_line / torch.clamp(torch.norm(epipolar_line[:, :2, :], dim=1, keepdim=True), min=1e-8)
        essential_cost = torch.abs(torch.sum(coord2_h * epipolar_line_, dim=1))  # Bxn
        return essential_cost

    def sg_loss(self,PT,scores):
        PT = PT.int()
        P_loss = scores[0,PT[:,0], PT[:,1]]
        loss = -P_loss.mean()
        return loss

    def set_weight(self, inverse_std, mask=None, regularizer=0.0):
        if self.config['use_std_as_weight']:
            # inverse_std = 1. / torch.clamp(std+regularizer, min=1e-10)
            weight = inverse_std / torch.mean(inverse_std)
            weight = weight.detach()  # Bxn
        else:
            weight = torch.ones_like(std)

        if mask is not None:
            weight *= mask.float()
            weight /= (torch.mean(weight) + 1e-8)
        return weight 

    def forward(self, inputs, outputs, processed):
        all_matches = inputs["matches"][0]
        scores = processed['scores']
        loss = self.sg_loss(all_matches, scores)

        return loss